#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
文档内容搜索工具 (CLI版本)
功能：在指定目录下搜索包含关键字的文件，支持多种文档格式
"""
import os
import sys
import argparse
import logging
import time
import datetime
import csv
import threading
import queue
import subprocess

# 导入文档处理相关库
try:
    from PyPDF2 import PdfReader
    has_pypdf2 = True
except ImportError:
    has_pypdf2 = False
    print("警告: PyPDF2库未安装，将无法处理PDF文件。")

try:
    import openpyxl
    has_openpyxl = True
except ImportError:
    has_openpyxl = False
    print("警告: openpyxl库未安装，将无法处理Excel文件。")

try:
    import python_pptx as pptx
    has_pptx = True
except ImportError:
    has_pptx = False
    print("警告: python-pptx库未安装，将无法处理PowerPoint文件。")

try:
    from docx import Document
    has_docx = True
except ImportError:
    has_docx = False
    print("警告: python-docx库未安装，将无法处理Word文件。")

# 如果缺少依赖，提示用户安装
missing_deps = []
if not has_pypdf2:
    missing_deps.append("PyPDF2")
if not has_openpyxl:
    missing_deps.append("openpyxl")
if not has_pptx:
    missing_deps.append("python-pptx")
if not has_docx:
    missing_deps.append("python-docx")

if missing_deps:
    print(f"\n请安装缺失的依赖库以获得完整功能:")
    print(f"pip install {' '.join(missing_deps)}")
    print(f"或使用: pip install -r requirements.txt\n")

# 配置日志系统
def setup_logger():
    """配置日志系统，确保中文显示正常"""
    logger = logging.getLogger('DocSearchCLI')
    logger.setLevel(logging.DEBUG)
    
    # 创建控制台处理器
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    
    # 创建文件处理器
    log_file = f"docsearch_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    file_handler = logging.FileHandler(log_file, encoding='utf-8')
    file_handler.setLevel(logging.DEBUG)
    
    # 定义日志格式
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(formatter)
    file_handler.setFormatter(formatter)
    
    # 添加处理器到logger
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)
    
    return logger

class DocSearchCLI:
    """文档内容搜索工具的命令行版本"""
    def __init__(self, directory, keyword, log_level=logging.INFO):
        self.directory = directory
        self.keyword = keyword
        self.logger = setup_logger()
        self.search_results = []
        self.search_cancelled = False
        self.total_files = 0
        self.processed_files = 0
        self.found_count = 0
        self.result_queue = queue.Queue()
        
    def start_search(self):
        """开始搜索过程"""
        # 验证输入参数
        if not os.path.isdir(self.directory):
            self.logger.error(f"目录 '{self.directory}' 不存在或不可访问")
            return False
            
        if not self.keyword:
            self.logger.error("搜索关键字不能为空")
            return False
            
        self.logger.info(f"开始搜索 - 目录: {self.directory}, 关键字: {self.keyword}")
        
        # 创建搜索线程
        search_thread = threading.Thread(target=self.search_in_thread)
        search_thread.daemon = True
        search_thread.start()
        
        # 处理搜索结果
        self.process_search_results()
        
        return True
        
    def search_in_thread(self):
        """在后台线程中执行搜索"""
        try:
            # 发送初始状态
            self.result_queue.put(('status', "正在计算文件总数..."))
            
            # 计算总文件数
            for root, _, files in os.walk(self.directory):
                self.total_files += len(files)
            
            self.logger.debug(f"总共找到 {self.total_files} 个文件")
            
            # 搜索文件
            for root, _, files in os.walk(self.directory):
                if self.search_cancelled:
                    self.result_queue.put(('cancelled', None))
                    return
                    
                for file in files:
                    if self.search_cancelled:
                        self.result_queue.put(('cancelled', None))
                        return
                        
                    self.processed_files += 1
                    file_path = os.path.join(root, file)
                    
                    # 发送状态更新 - 每处理10个文件才更新一次状态，避免更新太频繁
                    if self.processed_files % 10 == 0 or self.processed_files == 1:
                        progress = f"{self.processed_files/100:.0f}%" if self.total_files > 0 else "处理中"
                        self.result_queue.put(('status', f"搜索中... {self.processed_files}/{self.total_files} 文件 ({progress})"))
                        
                    # 检查文件名
                    if self.keyword.lower() in file.lower():
                        file_size = os.path.getsize(file_path)
                        modified_time = os.path.getmtime(file_path)
                        self.search_results.append({
                            'path': file_path,
                            'name': file,
                            'size': file_size,
                            'modified_time': modified_time,
                            'match_type': '文件名'
                        })
                        self.found_count += 1
                        self.logger.debug(f"找到匹配文件名: {file_path}")
                        continue
                        
                    # 检查文件内容
                    try:
                        content = self.extract_text(file_path)
                        if content and self.keyword.lower() in content.lower():
                            file_size = os.path.getsize(file_path)
                            modified_time = os.path.getmtime(file_path)
                            self.search_results.append({
                                'path': file_path,
                                'name': file,
                                'size': file_size,
                                'modified_time': modified_time,
                                'match_type': '文件内容'
                            })
                            self.found_count += 1
                            self.logger.debug(f"找到匹配内容: {file_path}")
                    except Exception as e:
                        # 记录错误但继续搜索
                        error_msg = f"无法读取文件 {file_path}: {str(e)}"
                        self.result_queue.put(('error', error_msg))
                        self.logger.warning(error_msg)
            
            # 搜索完成
            self.result_queue.put(('finished', (self.found_count, self.processed_files)))
        except Exception as e:
            self.logger.error(f"搜索过程中发生错误: {str(e)}")
            self.result_queue.put(('error', str(e)))
    
    def extract_text(self, file_path):
        """根据文件类型提取文本内容"""
        _, ext = os.path.splitext(file_path)
        ext = ext.lower()
        
        content = ""
        
        try:
            if ext == '.txt':
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    content = f.read()
            elif ext == '.docx' and has_docx:
                doc = Document(file_path)
                content = "\n".join([para.text for para in doc.paragraphs])
            elif ext == '.pdf' and has_pypdf2:
                with open(file_path, 'rb') as f:
                    reader = PdfReader(f)
                    for page in reader.pages:
                        page_text = page.extract_text() or ""
                        content += page_text + "\n"
            elif ext == '.xlsx' and has_openpyxl:
                wb = openpyxl.load_workbook(file_path, read_only=True)
                for sheet_name in wb.sheetnames:
                    sheet = wb[sheet_name]
                    for row in sheet.iter_rows():
                        for cell in row:
                            if cell.value:
                                content += str(cell.value) + "\t"
                        content += "\n"
            elif ext == '.pptx' and has_pptx:
                prs = pptx.Presentation(file_path)
                for slide in prs.slides:
                    for shape in slide.shapes:
                        if hasattr(shape, 'text'):
                            content += shape.text + "\n"
        except Exception as e:
            raise e
            
        return content
    
    def process_search_results(self):
        """处理搜索线程发送的结果"""
        while True:
            try:
                # 阻塞方式获取队列中的消息
                message_type, data = self.result_queue.get(timeout=0.1)
                
                if message_type == 'status':
                    # 更新状态
                    self.logger.info(data)
                elif message_type == 'error':
                    # 显示错误信息
                    self.logger.warning(data)
                elif message_type == 'finished':
                    # 搜索完成
                    found_count, processed_files = data
                    self.logger.info(f"搜索完成。共处理 {processed_files} 个文件，找到 {found_count} 个匹配文件。")
                    
                    # 显示结果
                    self.display_results()
                    return
                elif message_type == 'cancelled':
                    # 搜索被取消
                    self.logger.info("搜索已取消。")
                    return
                    
                # 标记任务完成
                self.result_queue.task_done()
            except queue.Empty:
                # 队列为空，继续等待
                continue
            except KeyboardInterrupt:
                # 捕获用户中断
                self.cancel_search()
                self.logger.info("搜索被用户中断。")
                return
    
    def display_results(self):
        """显示搜索结果"""
        if not self.search_results:
            self.logger.info("未找到匹配的文件。")
            return
            
        # 按文件名排序结果
        sorted_results = sorted(self.search_results, key=lambda x: x['name'].lower())
        
        self.logger.info(f"\n===== 搜索结果 (共 {len(sorted_results)} 个) =====")
        
        for i, result in enumerate(sorted_results, 1):
            # 获取修改时间的可读格式
            modified_str = datetime.datetime.fromtimestamp(result['modified_time']).strftime('%Y-%m-%d %H:%M:%S')
            
            # 获取文件大小的可读格式
            size_str = self.get_readable_size(result['size'])
            
            # 输出结果信息
            self.logger.info(f"{i}. {result['path']}")
            self.logger.info(f"   匹配: {result['match_type']}, 大小: {size_str}, 修改时间: {modified_str}")
            
    def get_readable_size(self, size_bytes):
        """将字节大小转换为人类可读的格式"""
        if size_bytes < 1024:
            return f"{size_bytes} B"
        elif size_bytes < 1024 * 1024:
            return f"{size_bytes / 1024:.2f} KB"
        elif size_bytes < 1024 * 1024 * 1024:
            return f"{size_bytes / (1024 * 1024):.2f} MB"
        else:
            return f"{size_bytes / (1024 * 1024 * 1024):.2f} GB"
            
    def cancel_search(self):
        """取消搜索过程"""
        self.search_cancelled = True
        
    def export_results(self, export_format='csv', output_file=None):
        """导出搜索结果"""
        if not self.search_results:
            self.logger.error("没有可导出的搜索结果")
            return False
            
        # 确定导出格式和文件扩展名
        if export_format.lower() == 'csv':
            file_ext = '.csv'
            file_types = ["CSV文件", "*.csv"]
        elif export_format.lower() == 'txt':
            file_ext = '.txt'
            file_types = ["文本文件", "*.txt"]
        else:
            self.logger.error(f"不支持的导出格式: {export_format}")
            return False
            
        # 确定输出文件路径
        if not output_file:
            current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = f"搜索结果_{current_time}{file_ext}"
            
        try:
            # 按文件名排序结果
            sorted_results = sorted(self.search_results, key=lambda x: x['name'].lower())
            
            if export_format.lower() == 'csv':
                self.export_to_csv(sorted_results, output_file)
            else:
                self.export_to_txt(sorted_results, output_file)
                
            self.logger.info(f"搜索结果已成功导出到 {output_file}")
            return True
        except Exception as e:
            self.logger.error(f"导出失败: {str(e)}")
            return False
            
    def export_to_csv(self, results, file_path):
        """导出搜索结果到CSV文件"""
        with open(file_path, 'w', newline='', encoding='utf-8-sig') as csvfile:
            fieldnames = ['序号', '文件名', '文件路径', '匹配类型', '大小', '修改时间']
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            
            # 写入表头
            writer.writeheader()
            
            # 写入数据
            for i, result in enumerate(results, 1):
                modified_str = datetime.datetime.fromtimestamp(result['modified_time']).strftime('%Y-%m-%d %H:%M:%S')
                size_str = self.get_readable_size(result['size'])
                
                writer.writerow({
                    '序号': i,
                    '文件名': result['name'],
                    '文件路径': result['path'],
                    '匹配类型': result['match_type'],
                    '大小': size_str,
                    '修改时间': modified_str
                })
                
    def export_to_txt(self, results, file_path):
        """导出搜索结果到文本文件"""
        with open(file_path, 'w', encoding='utf-8') as txtfile:
            # 写入导出信息
            txtfile.write(f"文档内容搜索工具 - 搜索结果导出\n")
            txtfile.write(f"导出时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            txtfile.write(f"搜索目录: {self.directory}\n")
            txtfile.write(f"搜索关键字: {self.keyword}\n")
            txtfile.write(f"结果数量: {len(results)}\n")
            txtfile.write(f"排序方式: 按文件名\n")
            txtfile.write("=" * 80 + "\n\n")
            
            # 写入结果数据
            for i, result in enumerate(results, 1):
                modified_str = datetime.datetime.fromtimestamp(result['modified_time']).strftime('%Y-%m-%d %H:%M:%S')
                size_str = self.get_readable_size(result['size'])
                
                txtfile.write(f"{i}. 文件名: {result['name']}\n")
                txtfile.write(f"   路径: {result['path']}\n")
                txtfile.write(f"   匹配: {result['match_type']}\n")
                txtfile.write(f"   大小: {size_str}\n")
                txtfile.write(f"   修改时间: {modified_str}\n")
                txtfile.write("-" * 80 + "\n")
                
    def open_file_location(self, file_path):
        """打开文件所在的文件夹"""
        if not file_path or not os.path.exists(file_path):
            self.logger.error(f"文件不存在: {file_path}")
            return False
            
        try:
            # 获取文件所在目录
            directory = os.path.dirname(file_path)
            
            # 在Windows上打开目录并选中文件
            if os.name == 'nt':  # Windows系统
                # 使用 explorer /select 命令打开目录并选中文件
                subprocess.Popen(f'explorer /select,"{file_path}"')
            else:
                # 其他系统仅打开目录
                subprocess.Popen(['xdg-open', directory])
            
            self.logger.info(f"已打开文件位置: {directory}")
            return True
        except Exception as e:
            self.logger.error(f"无法打开文件位置: {str(e)}")
            return False

def main():
    """主函数，处理命令行参数并启动搜索"""
    parser = argparse.ArgumentParser(description='文档内容搜索工具 (CLI版本)')
    parser.add_argument('directory', help='要搜索的目录路径')
    parser.add_argument('keyword', help='搜索关键字')
    parser.add_argument('-e', '--export', choices=['csv', 'txt'], help='导出格式')
    parser.add_argument('-o', '--output', help='导出文件路径')
    parser.add_argument('-v', '--verbose', action='store_true', help='显示详细日志信息')
    
    args = parser.parse_args()
    
    # 设置日志级别
    log_level = logging.DEBUG if args.verbose else logging.INFO
    
    # 创建搜索实例
    search_tool = DocSearchCLI(args.directory, args.keyword, log_level)
    
    try:
        # 开始搜索
        success = search_tool.start_search()
        
        # 如果指定了导出，则导出结果
        if success and args.export:
            search_tool.export_results(args.export, args.output)
            
    except KeyboardInterrupt:
        search_tool.cancel_search()
        print("\n搜索已取消。")

if __name__ == "__main__":
    main()